package william.heap;

/**
 * @author ZhangShenao
 * @date 2024/2/20
 * @description <a href="https://leetcode.cn/problems/xx4gT2/description/">...</a>
 */
public class LCR_076_数组中的第K大元素 {
    /**
     * 采用堆排序
     * 自顶向下建立大根堆shiftUp操作
     * 将堆顶元素（数组0位置元素）与数组N-1位置元素交换,此时N-1位置元素即为第1大的元素
     * 对数组0位置元素进行shiftDown操作,维持大根堆
     * 上述过程持续k-1次
     * 最终堆顶元素（数组0位置元素）即为第K大元素,直接返回。
     */
    public int findKthLargest(int[] nums, int k) {
        //边界条件校验
        if (nums == null || nums.length < k) {
            throw new IllegalArgumentException("invalid input!");
        }

        //将整个数组构造成大根堆
        buildMaxHeap(nums);
        int heapSize = nums.length;

        //进行k-1次元素交换和shiftDown操作
        for (int i = k; i > 1; i--) {
            swap(nums, 0, --heapSize);
            shiftDown(nums, heapSize, 0);
        }

        //最终堆顶元素（数组0位置元素）即为第K大元素,直接返回。
        return nums[0];
    }

    /**
     * 将数组构造成大根堆
     */
    private void buildMaxHeap(int[] nums) {
        //从叶子节点起,依次进行shiftDown操作
        for (int i = 0; i < nums.length; i++) {
            shiftUp(nums, i);
        }
    }

    /**
     * 上浮操作
     * 针对每个节点,如果比父节点大,则与父节点进行交换
     * 持续该过程,直到比父节点小,或到达根节点
     */
    private void shiftUp(int[] nums, int i) {
        while (i > 0 && nums[i] > nums[(i - 1) / 2]) {
            swap(nums, i, (i - 1) / 2);
            i = (i - 1) / 2;
        }
    }

    /**
     * 下沉操作
     * 针对每个节点,找到其左、右子节点中的最大节点,与其进行交换
     */
    private void shiftDown(int[] nums, int heapSize, int idx) {
        while (idx < heapSize) {
            //找到当前节点的左、右子节点中的最大节点
            int leftIdx = idx * 2 + 1;
            if (leftIdx >= heapSize) {    //无左子节点,终止操作
                break;
            }
            int maxIdx = leftIdx;
            int rightIdx = idx * 2 + 2;
            if (rightIdx < heapSize && nums[rightIdx] > nums[leftIdx]) {
                maxIdx = rightIdx;
            }

            //当前节点比左、右子节点都大,终止操作
            if (nums[idx] >= nums[maxIdx]) {
                break;
            }

            //将当前节点与最大子节点进行交换,并持续该过程
            swap(nums, idx, maxIdx);
            idx = maxIdx;
        }
    }

    private void swap(int[] nums, int i, int j) {
        int tmp = nums[i];
        nums[i] = nums[j];
        nums[j] = tmp;
    }
}
